//
// Created by twetec on 25-8-13.
//

#include "handle_postgresql_probe_task.h"

#include <postgresql/libpq-fe.h>
#include <postgresql/pg_config_ext.h>
#include <postgresql/postgres_ext.h>
#include <future>

#include <bs_thread_pool/BS_thread_pool.hpp>

using namespace std::chrono_literals;

namespace network_prober::probe
{
// 线程池
static BS::thread_pool s_iopool{1000};

// PostgreSQLProbeIPv4::PostgreSQLProbeIPv4(
//     std::string_view interface,
//     std::string_view target,
//     std::uint32_t port,
//     std::string_view userName,
//     std::string_view password,
//     std::string_view databaseName,
//     std::string_view queryStatement
// )
// : m_interface(interface), m_target(target), m_port(port), m_userName(userName),
// m_password(password), m_databaseName(databaseName), m_queryStatement(queryStatement)
// {
//
// }
//
// void PostgreSQLProbeIPv4::setOnProbedCallback(std::function<void(std::int64_t)> callback)
// {
//     m_onProbedCallback = std::move(callback);
// }
//
// void PostgreSQLProbeIPv4::probe()
// {
//     // 建立连接时，开始计时
//     m_probeStartTime = std::chrono::steady_clock::now();
//
//     // 指向当前对象的shared_ptr
//     auto self{shared_from_this()};
//
//     // 发起探测
//     m_onProbedCallbackFuture = std::async(std::launch::async, [this, self] {
//         // 初始化数据库连接配置
//         const char* paramKeys[] = {
//             "host", // 要连接的主机的数字 IP 地址。 这应该是标准 IPv4 地址格式，例如 172.28.40.9。 如果您的机器支持 IPv6，您也可以使用这些地址
//             "port", // 服务器主机上连接的端口号
//             "dbname",   // 数据库名称
//             "user", // 要连接的 PostgreSQL 用户名
//             "password", // 服务器要求密码验证时使用的密码
//             "connect_timeout",   // 连接时等待的最长时间，以秒为单位（写为十进制整数，例如 10）。 零、负数或未指定意味着无限期等待。 允许的最小超时为 2 秒，因此值 1 被解释为 2
//             NULL
//         };
//         const char* paramValues[] = {
//             m_target.c_str(),
//             std::to_string(m_port).c_str(),
//             m_databaseName.c_str(),
//             m_userName.c_str(),
//             m_password.c_str(),
//             "13",
//             NULL
//         };
//
//         // 连接PQ数据库服务器
//         PGconn* connPtr = PQconnectdbParams(paramKeys, paramValues, 0);
//
//         // 检查连接是否成功
//         ConnStatusType status = PQstatus(connPtr);
//         // 如果连接状态不是成功
//         if (status != CONNECTION_OK) {
//            // 错误信息
//            SPDLOG_INFO("PgSQL观测目标连接失败：{}", PQerrorMessage(connPtr));
//            // 关闭与服务器的而连接，同时释放了PGconn对象的内存 注：即使服务器连接失败，也要调用finish函数，回收内存
//            PQfinish(connPtr);
//            // 回调结果
//            m_onProbedCallback(-1);
//            // 结束执行
//            return;
//         }
//
//         // 构建SQL语句
//         const char* sql = m_queryStatement.empty()? "SELECT 1" : m_queryStatement.c_str();
//
//         // 让服务器执行SQL语句
//         PGresult* execResult = PQexec(connPtr, sql);
//
//         // 检查查询是否成功
//         ExecStatusType execStatus = PQresultStatus(execResult);
//         // 如果查询语句执行失败
//         if (execStatus != PGRES_TUPLES_OK) {
//             // 错误信息
//             SPDLOG_INFO("PgSQL观测目标查询失败：{}", PQerrorMessage(connPtr));
//             // 释放PGresult查询结果，有PGresult对象，就必须释放
//             PQclear(execResult);
//             // 关闭与服务器的而连接，同时释放了PGconn对象的内存 注：即使服务器连接失败，也要调用finish函数，回收内存
//             PQfinish(connPtr);
//             // 回调结果
//             m_onProbedCallback(-1);
//             // 结束执行
//             return;
//         }
//
//         // 释放PGresult查询结果，有PGresult对象，就必须释放
//         PQclear(execResult);
//         // 关闭与服务器的而连接，同时释放了PGconn对象的内存 注：即使服务器连接失败，也要调用finish函数，回收内存
//         PQfinish(connPtr);
//
//         // 回调结果
//         m_onProbedCallback(std::chrono::duration_cast<std::chrono::nanoseconds>(std::chrono::steady_clock::now() - m_probeStartTime).count());
//     });
// }
//
// PostgreSQLProbeIPv4::~PostgreSQLProbeIPv4()
// {
//
// }

HandlePostgreSQLProbeIPv4Awaiter::HandlePostgreSQLProbeIPv4Awaiter(const database::ProbeTaskModel &probeTask)
: m_interface(probeTask.interface), m_target(probeTask.address), m_port(probeTask.port), m_userName(probeTask.userName),
m_password(probeTask.password), m_databaseName(probeTask.databaseName), m_queryStatement(probeTask.queryStatement)
{
}

bool HandlePostgreSQLProbeIPv4Awaiter::await_ready()
{
    return false;
}

void HandlePostgreSQLProbeIPv4Awaiter::await_suspend(std::coroutine_handle<> handle)
{
    s_iopool.submit_task([this, handle] {
        // 线程锁
        std::mutex mtx;
        std::condition_variable cv;

        // 请求计时 - 开始时间
        std::chrono::time_point<std::chrono::steady_clock> probeStartTime = std::chrono::steady_clock::now();

        // 发起探测
        std::future future = std::async(std::launch::async, [this, &probeStartTime, &mtx, &cv] {
            // 初始化数据库连接配置
            const char* paramKeys[] = {
                "host", // 要连接的主机的数字 IP 地址。 这应该是标准 IPv4 地址格式，例如 172.28.40.9。 如果您的机器支持 IPv6，您也可以使用这些地址
                "port", // 服务器主机上连接的端口号
                "dbname",   // 数据库名称
                "user", // 要连接的 PostgreSQL 用户名
                "password", // 服务器要求密码验证时使用的密码
                "connect_timeout",   // 连接时等待的最长时间，以秒为单位（写为十进制整数，例如 10）。 零、负数或未指定意味着无限期等待。 允许的最小超时为 2 秒，因此值 1 被解释为 2
                NULL
            };
            const char* paramValues[] = {
                m_target.c_str(),
                std::to_string(m_port).c_str(),
                m_databaseName.c_str(),
                m_userName.c_str(),
                m_password.c_str(),
                "6",
                NULL
            };

            // 连接PQ数据库服务器
            PGconn* connPtr = PQconnectdbParams(paramKeys, paramValues, 0);

            // 检查连接是否成功
            ConnStatusType status = PQstatus(connPtr);
            // 如果连接状态不是成功
            if (status != CONNECTION_OK) {
                // 错误信息
                SPDLOG_INFO("PgSQL观测目标连接失败：{}", PQerrorMessage(connPtr));
                // 关闭与服务器的而连接，同时释放了PGconn对象的内存 注：即使服务器连接失败，也要调用finish函数，回收内存
                PQfinish(connPtr);
                // 加锁
                std::unique_lock<std::mutex> lock(mtx);
                // 保存延迟
                m_delay = -1;
                // 通知延迟保存成功
                cv.notify_one();
                // 结束执行
                return;
            }

            // 构建SQL语句
            const char* sql = m_queryStatement.empty()? "SELECT 1" : m_queryStatement.c_str();

            // 让服务器执行SQL语句
            PGresult* execResult = PQexec(connPtr, sql);

            // 检查查询是否成功
            ExecStatusType execStatus = PQresultStatus(execResult);
            // 如果查询语句执行失败
            if (execStatus != PGRES_TUPLES_OK) {
                // 错误信息
                SPDLOG_INFO("PgSQL观测目标查询失败：{}", PQerrorMessage(connPtr));
                // 释放PGresult查询结果，有PGresult对象，就必须释放
                PQclear(execResult);
                // 关闭与服务器的而连接，同时释放了PGconn对象的内存 注：即使服务器连接失败，也要调用finish函数，回收内存
                PQfinish(connPtr);
                // 加锁
                std::unique_lock<std::mutex> lock(mtx);
                // 保存延迟
                m_delay = -1;
                // 通知延迟保存成功
                cv.notify_one();
                // 结束执行
                return;
            }

            // 释放PGresult查询结果，有PGresult对象，就必须释放
            PQclear(execResult);
            // 关闭与服务器的而连接，同时释放了PGconn对象的内存 注：即使服务器连接失败，也要调用finish函数，回收内存
            PQfinish(connPtr);

            // 加锁
            std::unique_lock<std::mutex> lock(mtx);
            // 保存延迟
            m_delay = std::chrono::duration_cast<std::chrono::nanoseconds>(std::chrono::steady_clock::now() - probeStartTime).count();
            // 通知延迟保存成功
            cv.notify_one();
        });

        // 等待响应（12s）
        std::unique_lock<std::mutex> lock(mtx);
        std::cv_status cvStatus = cv.wait_for(lock, 8s);

        // 如果等待超时,则将延迟置为-1
        if (cvStatus == std::cv_status::timeout) m_delay = -1;

        // 恢复暂停的协程
        handle.resume();
    });
}

std::int64_t HandlePostgreSQLProbeIPv4Awaiter::await_resume()
{
    return m_delay;
}


// PostgreSQLProbeIPv6::PostgreSQLProbeIPv6(
//     std::string_view interface,
//     std::string_view target,
//     std::uint32_t port,
//     std::string_view userName,
//     std::string_view password,
//     std::string_view databaseName,
//     std::string_view queryStatement
// )
// : m_interface(interface), m_target(target), m_port(port), m_userName(userName),
// m_password(password), m_databaseName(databaseName), m_queryStatement(queryStatement)
// {
//
// }
//
// void PostgreSQLProbeIPv6::setOnProbedCallback(std::function<void(std::int64_t)> callback)
// {
//     m_onProbedCallback = std::move(callback);
// }
//
// void PostgreSQLProbeIPv6::probe()
// {
//     // 建立连接时，开始计时
//     m_probeStartTime = std::chrono::steady_clock::now();
//
//     // 指向当前对象的shared_ptr
//     auto self{shared_from_this()};
//
//     // 发起探测
//     m_onProbedCallbackFuture = std::async(std::launch::async, [this, self] {
//         // 初始化数据库连接配置
//         const char* paramKeys[] = {
//             "host", // 要连接的主机的数字 IP 地址。 这应该是标准 IPv4 地址格式，例如 172.28.40.9。 如果您的机器支持 IPv6，您也可以使用这些地址
//             "port", // 服务器主机上连接的端口号
//             "dbname",   // 数据库名称
//             "user", // 要连接的 PostgreSQL 用户名
//             "password", // 服务器要求密码验证时使用的密码
//             "connect_timeout",   // 连接时等待的最长时间，以秒为单位（写为十进制整数，例如 10）。 零、负数或未指定意味着无限期等待。 允许的最小超时为 2 秒，因此值 1 被解释为 2
//             NULL
//         };
//         const char* paramValues[] = {
//             m_target.c_str(),
//             std::to_string(m_port).c_str(),
//             m_databaseName.c_str(),
//             m_userName.c_str(),
//             m_password.c_str(),
//             "13",
//             NULL
//         };
//
//         // 连接PQ数据库服务器
//         PGconn* connPtr = PQconnectdbParams(paramKeys, paramValues, 0);
//
//         // 检查连接是否成功
//         ConnStatusType status = PQstatus(connPtr);
//         // 如果连接状态不是成功
//         if (status != CONNECTION_OK) {
//            // 错误信息
//            SPDLOG_INFO("PgSQL观测目标连接失败：{}", PQerrorMessage(connPtr));
//            // 关闭与服务器的而连接，同时释放了PGconn对象的内存 注：即使服务器连接失败，也要调用finish函数，回收内存
//            PQfinish(connPtr);
//            // 回调结果
//            m_onProbedCallback(-1);
//            // 结束执行
//            return;
//         }
//
//         // 构建SQL语句
//         const char* sql = m_queryStatement.empty()? "SELECT 1" : m_queryStatement.c_str();
//
//         // 让服务器执行SQL语句
//         PGresult* execResult = PQexec(connPtr, sql);
//
//         // 检查查询是否成功
//         ExecStatusType execStatus = PQresultStatus(execResult);
//         // 如果查询语句执行失败
//         if (execStatus != PGRES_TUPLES_OK) {
//             // 错误信息
//             SPDLOG_INFO("PgSQL观测目标查询失败：{}", PQerrorMessage(connPtr));
//             // 释放PGresult查询结果，有PGresult对象，就必须释放
//             PQclear(execResult);
//             // 关闭与服务器的而连接，同时释放了PGconn对象的内存 注：即使服务器连接失败，也要调用finish函数，回收内存
//             PQfinish(connPtr);
//             // 回调结果
//             m_onProbedCallback(-1);
//             // 结束执行
//             return;
//         }
//
//         // 释放PGresult查询结果，有PGresult对象，就必须释放
//         PQclear(execResult);
//         // 关闭与服务器的而连接，同时释放了PGconn对象的内存 注：即使服务器连接失败，也要调用finish函数，回收内存
//         PQfinish(connPtr);
//
//         // 回调结果
//         m_onProbedCallback(std::chrono::duration_cast<std::chrono::nanoseconds>(std::chrono::steady_clock::now() - m_probeStartTime).count());
//     });
// }
//
// PostgreSQLProbeIPv6::~PostgreSQLProbeIPv6()
// {
//
// }

HandlePostgreSQLProbeIPv6Awaiter::HandlePostgreSQLProbeIPv6Awaiter(const database::ProbeTaskModel &probeTask)
: m_interface(probeTask.interface), m_target(probeTask.address), m_port(probeTask.port), m_userName(probeTask.userName),
m_password(probeTask.password), m_databaseName(probeTask.databaseName), m_queryStatement(probeTask.queryStatement)
{
}

bool HandlePostgreSQLProbeIPv6Awaiter::await_ready()
{
    return false;
}

void HandlePostgreSQLProbeIPv6Awaiter::await_suspend(std::coroutine_handle<> handle)
{
    s_iopool.submit_task([this, handle] {
        // 线程锁
        std::mutex mtx;
        std::condition_variable cv;

        // 请求计时 - 开始时间
        std::chrono::time_point<std::chrono::steady_clock> probeStartTime = std::chrono::steady_clock::now();

        // 发起探测
        std::future future = std::async(std::launch::async, [this, &probeStartTime, &mtx, &cv] {
            // 初始化数据库连接配置
            const char*paramKeys[] = {
                "host",            // 要连接的主机的数字 IP 地址。 这应该是标准 IPv4 地址格式，例如 172.28.40.9。 如果您的机器支持 IPv6，您也可以使用这些地址
                "port",            // 服务器主机上连接的端口号
                "dbname",          // 数据库名称
                "user",            // 要连接的 PostgreSQL 用户名
                "password",        // 服务器要求密码验证时使用的密码
                "connect_timeout", // 连接时等待的最长时间，以秒为单位（写为十进制整数，例如 10）。 零、负数或未指定意味着无限期等待。 允许的最小超时为 2 秒，因此值 1 被解释为 2
                NULL
            };
            const char *paramValues[] = {
                m_target.c_str(),
                std::to_string(m_port).c_str(),
                m_databaseName.c_str(),
                m_userName.c_str(),
                m_password.c_str(),
                "6",
                NULL
            };

            // 连接PQ数据库服务器
            PGconn *connPtr = PQconnectdbParams(paramKeys, paramValues, 0);

            // 检查连接是否成功
            ConnStatusType status = PQstatus(connPtr);
            // 如果连接状态不是成功
            if (status != CONNECTION_OK) {
                // 错误信息
                SPDLOG_INFO("PgSQL观测目标连接失败：{}", PQerrorMessage(connPtr));
                // 关闭与服务器的而连接，同时释放了PGconn对象的内存 注：即使服务器连接失败，也要调用finish函数，回收内存
                PQfinish(connPtr);
                // 加锁
                std::unique_lock<std::mutex> lock(mtx);
                // 保存延迟
                m_delay = -1;
                // 通知延迟保存成功
                cv.notify_one();
                // 结束执行
                return;
            }

            // 构建SQL语句
            const char *sql = m_queryStatement.empty() ? "SELECT 1" : m_queryStatement.c_str();

            // 让服务器执行SQL语句
            PGresult *execResult = PQexec(connPtr, sql);

            // 检查查询是否成功
            ExecStatusType execStatus = PQresultStatus(execResult);
            // 如果查询语句执行失败
            if (execStatus != PGRES_TUPLES_OK) {
                // 错误信息
                SPDLOG_INFO("PgSQL观测目标查询失败：{}", PQerrorMessage(connPtr));
                // 释放PGresult查询结果，有PGresult对象，就必须释放
                PQclear(execResult);
                // 关闭与服务器的而连接，同时释放了PGconn对象的内存 注：即使服务器连接失败，也要调用finish函数，回收内存
                PQfinish(connPtr);
                // 加锁
                std::unique_lock<std::mutex> lock(mtx);
                // 保存延迟
                m_delay = -1;
                // 通知延迟保存成功
                cv.notify_one();
                // 结束执行
                return;
            }

            // 释放PGresult查询结果，有PGresult对象，就必须释放
            PQclear(execResult);
            // 关闭与服务器的而连接，同时释放了PGconn对象的内存 注：即使服务器连接失败，也要调用finish函数，回收内存
            PQfinish(connPtr);

            // 加锁
            std::unique_lock<std::mutex> lock(mtx);
            // 保存延迟
            m_delay = std::chrono::duration_cast<std::chrono::nanoseconds>(std::chrono::steady_clock::now() - probeStartTime).count();
            // 通知延迟保存成功
            cv.notify_one();
        });

        // 等待响应（12s）
        std::unique_lock<std::mutex> lock(mtx);
        std::cv_status cvStatus = cv.wait_for(lock, 8s);

        // 如果等待超时,则将延迟置为-1
        if (cvStatus == std::cv_status::timeout) m_delay = -1;

        // 恢复暂停的协程
        handle.resume();
    });
}

std::int64_t HandlePostgreSQLProbeIPv6Awaiter::await_resume()
{
    return m_delay;
}
}